package edu.nepu.flink.api.sink;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.serialization.SimpleStringEncoder;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.connector.datagen.source.DataGeneratorSource;
import org.apache.flink.connector.datagen.source.GeneratorFunction;
import org.apache.flink.connector.file.sink.FileSink;
import org.apache.flink.core.fs.Path;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.filesystem.OutputFileConfig;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.DateTimeBucketAssigner;
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.DefaultRollingPolicy;
import org.apache.flink.streaming.api.functions.source.datagen.DataGenerator;

import java.time.Duration;
import java.time.ZoneId;
import java.util.Random;
import java.util.UUID;

/**
 * @Date 2024/2/29 9:05
 * @Created by chenshuaijun
 */
public class SinkToFile {


    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(1);

        // 数据输出到文件，需要开启checkpoint，要不然文件会一直是写入的状态，外部无法访问
        env.enableCheckpointing(2000, CheckpointingMode.EXACTLY_ONCE);

        DataGeneratorSource<String> source = new DataGeneratorSource<>(new GeneratorFunction<Long, String>() {
            Random random = new Random();

            @Override
            public String map(Long aLong) throws Exception {
                return UUID.randomUUID().toString().substring(0, 4) + "," + String.valueOf(random.nextInt(100));
            }
        },
                10000,
                RateLimiterStrategy.perSecond(100),
                TypeInformation.of(String.class)
        );


        DataStreamSource<String> streamSource = env.fromSource(source, WatermarkStrategy.noWatermarks(),"date-gen");

        FileSink<String> fileSink = FileSink.<String>forRowFormat(
                        new Path("output"),
                        new SimpleStringEncoder<>()).
                withBucketAssigner(new DateTimeBucketAssigner<>("yyyy-MM-dd HH", ZoneId.systemDefault())).
                withOutputFileConfig(new OutputFileConfig("nepu", ".log")).
                withRollingPolicy(DefaultRollingPolicy.builder().
                        withMaxPartSize(new MemorySize(1024)).withRolloverInterval(Duration.ofSeconds(10)).build()
                ).build();


        streamSource.sinkTo(fileSink);
        env.execute();


    }
}
